from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from typing import Literal, Dict, Any, AsyncGenerator, Optional
from openai import OpenAI
import logging
import uvicorn
from config import settings


app = FastAPI()

class TwoDElementExtractionAPI:
    def __init__(self):
        self.client = OpenAI(
            api_key=settings.deepseek_api_key,
            base_url=settings.deepseek_base_url,
        )
        
        self.system_prompt = '''
<身份>
你是一名擅长中学数学的老师，你需要根据上传的数学题目，提取出题目和问题中包含的所有图形（圆、四边形、三角形等）、点、角、直线（对称轴、切线等）、线段和函数并以规定格式列出。
若题目中元素的信息在问题中需要求解得出（点的坐标、函数的解析式等），你必须计算答案并与其他元素一起列出，以保证输出答案的精确。
</身份>

<需求>
•   题目和各个小问中提到的所有元素包括计算过程中涉及的辅助线都必须包含，以保证输出答案的精确。
•   若题目中需要分类讨论，相关元素的所有情况都必须列出，以保证输出答案的精准。
•   输出严格参照下方格式，不得添加任何额外内容。
</需求>

<输出格式>
### 1. **点**

### 2. **角**

### 3. **直线**

### 4. **线段**

### 5. **函数**

### 6. **图形**
</输出格式>
'''

        logging.basicConfig(
            filename="twod_element_extraction_api.log",
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s"
        )

    async def extract_elements(self, problem: str) -> Dict[str, str]:
        """
        Extract 2D geometry elements from a math problem
        """
        try:
            logging.info(f"Received input: {problem}")
            response = self.client.chat.completions.create(
                model=settings.deepseek_reasoner_model,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": problem},
                ],
                stream=False
            )
            reasoning_content = response.choices[0].message.reasoning_content
            elements = response.choices[0].message.content
            logging.info(f"Extracted 2D elements: {elements}")
            return {
                "reasoning": reasoning_content,
                "elements": elements
            }
        except Exception as e:
            logging.error(f"Error in extract_elements: {str(e)}")
            raise

    async def extract_elements_stream(self, problem: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Stream the 2D geometry element extraction process
        """
        try:
            logging.info(f"Received input for streaming: {problem[:100]}...")
            logging.info(f"Using model: {settings.deepseek_reasoner_model}")
            
            response = self.client.chat.completions.create(
                model=settings.deepseek_reasoner_model,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": problem},
                ],
                stream=True
            )
            
            reasoning_chunks = []
            element_chunks = []
            chunk_count = 0
            
            for chunk in response:
                chunk_count += 1
                delta = chunk.choices[0].delta
                
                if delta.reasoning_content:
                    reasoning_chunks.append(delta.reasoning_content)
                    logging.info(f"Reasoning chunk {chunk_count}: {delta.reasoning_content[:50]}...")
                    yield {"type": "reasoning", "content": delta.reasoning_content}
                    
                elif delta.content:
                    element_chunks.append(delta.content)
                    logging.info(f"Element chunk {chunk_count}: {delta.content[:50]}...")
                    yield {"type": "element", "content": delta.content}
            
            # 汇总结果
            total_reasoning = ''.join(reasoning_chunks)
            total_elements = ''.join(element_chunks)
            
            logging.info(f"Stream完成 - 总chunk数: {chunk_count}")
            logging.info(f"推理内容总长度: {len(total_reasoning)}")
            logging.info(f"元素内容总长度: {len(total_elements)}")
            logging.info(f"最终2D元素内容: {total_elements}")
            
            # 检查元素内容是否有效，如果无效则使用非流式接口作为fallback
            if not total_elements or total_elements.strip() in ["###", ""] or len(total_elements.strip()) < 10:
                logging.warning("流式接口返回的元素内容无效，使用非流式接口作为fallback")
                try:
                    fallback_result = await self.extract_elements(problem)
                    fallback_elements = fallback_result.get('elements', '')
                    if fallback_elements and len(fallback_elements.strip()) > 10:
                        logging.info(f"Fallback成功，元素内容长度: {len(fallback_elements)}")
                        yield {"type": "element", "content": fallback_elements}
                    else:
                        logging.error("Fallback也失败了")
                        yield {"type": "error", "content": "2D元素提取失败：流式和非流式接口都无法获取有效元素"}
                except Exception as fallback_error:
                    logging.error(f"Fallback失败: {fallback_error}")
                    yield {"type": "error", "content": f"2D元素提取失败: {fallback_error}"}
                    
        except Exception as e:
            logging.error(f"Error in extract_elements_stream: {str(e)}")
            yield {"type": "error", "content": str(e)}


class ProblemInput(BaseModel):
    problem: str

class ElementsOutput(BaseModel):
    reasoning: str
    elements: str

@app.post("/extract_elements", response_model=ElementsOutput)
async def extract_elements(input: ProblemInput):
    try:
        logging.info(f"Received input: {input.problem}")
        api = TwoDElementExtractionAPI()
        result = await api.extract_elements(input.problem)
        return JSONResponse(content=result, status_code=200)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/extract_elements_stream")
async def extract_elements_stream(input: ProblemInput):
    """流式处理接口 - 兼容性保留"""
    try:
        logging.info(f"Received input for streaming: {input.problem}")
        api = TwoDElementExtractionAPI()
        
        async def generate_stream():
            async for chunk in api.extract_elements_stream(input.problem):
                yield f"data: {chunk}\n\n"
        
        return StreamingResponse(generate_stream(), media_type="text/event-stream")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run("ds_element_extraction_api:app", host="0.0.0.0", port=8002, reload=False)